{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Qh1AB1VcOBKs"
      },
      "source": [
        "##### Copyright 2019 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "tVlmNBATOIIW"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "2EZBfZhwiihm"
      },
      "source": [
        "# Spherical Harmonics Approximation\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/spherical_harmonics_approximation.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/spherical_harmonics_approximation.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5aywNaXvX505"
      },
      "source": [
        "This Colab covers an advanced topic and hence focuses on providing controllable\n",
        "toy examples to form a high level understanding of Spherical Harmonics and their\n",
        "use for lighting rather than providing step by step details. For those\n",
        "interested, these details are nevertheless available in the code. A great\n",
        "resource to form a good understanding of Spherical Harmonics and their use for\n",
        "lighting is\n",
        "[Spherical Harmonics Lighting: the Gritty Details](http://silviojemma.com/public/papers/lighting/spherical-harmonic-lighting.pdf).\n",
        "\n",
        "This Colab demonstrates how to approximate functions defined over a sphere using\n",
        "Spherical Harmonics. These can be used to approximate lighting and\n",
        "[reflectance](https://colab.research.google.com/github/tensorflow/graphics/blob/master/tensorflow_graphics/notebooks/reflectance.ipynb),\n",
        "leading to very efficient rendering.\n",
        "\n",
        "In more details, the following cells demonstrate:\n",
        "\n",
        "*   Approximation of lighting environments with Spherical Harmonics (SH)\n",
        "*   Approximation of the Lambertian BRDF with Zonal Harmonics (ZH)\n",
        "*   Rotation of Zonal Harmonics\n",
        "*   Rendering via Spherical Harmonics convolution of the SH lighting and ZH BRDF"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IT_MsDnVOLq1"
      },
      "source": [
        "## Setup \u0026 Imports\n",
        "If Tensorflow Graphics is not installed on your system, the following cell can install the Tensorflow Graphics package for you."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "lUMmhgr1OOuv"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_graphics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LW9kX4R1OU7K"
      },
      "source": [
        "Now that Tensorflow Graphics is installed, let's import everything needed to run the demo contained in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5qStA7s4F_FD"
      },
      "outputs": [],
      "source": [
        "###########\n",
        "# Imports #\n",
        "###########\n",
        "import math\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow_graphics.geometry.representation import grid\n",
        "from tensorflow_graphics.geometry.representation import ray\n",
        "from tensorflow_graphics.geometry.representation import vector\n",
        "from tensorflow_graphics.rendering.camera import orthographic\n",
        "from tensorflow_graphics.math import spherical_harmonics\n",
        "from tensorflow_graphics.math import math_helpers as tf_math\n",
        "\n",
        "tf.compat.v1.enable_eager_execution()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9Rt56g-nUiZZ"
      },
      "source": [
        "## Approximation of lighting with Spherical Harmonics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "_adaog1rOf8W"
      },
      "outputs": [],
      "source": [
        "#@title Controls { vertical-output: false, run: \"auto\" }\n",
        "max_band = 2  #@param { type: \"slider\", min: 0, max: 10 , step: 1 }\n",
        "\n",
        "#########################################################################\n",
        "# This cell creates a lighting function which we approximate with an SH #\n",
        "#########################################################################\n",
        "\n",
        "def image_to_spherical_coordinates(image_width, image_height):\n",
        "  pixel_grid_start = np.array((0, 0), dtype=type)\n",
        "  pixel_grid_end = np.array((image_width - 1, image_height - 1), dtype=type)\n",
        "  pixel_nb = np.array((image_width, image_height))\n",
        "  pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "  normalized_pixels = pixels / (image_width - 1, image_height - 1)\n",
        "  spherical_coordinates = tf_math.square_to_spherical_coordinates(\n",
        "      normalized_pixels)\n",
        "  return spherical_coordinates\n",
        "\n",
        "\n",
        "def light_function(theta, phi):\n",
        "  theta = tf.convert_to_tensor(theta)\n",
        "  phi = tf.convert_to_tensor(phi)\n",
        "  zero = tf.zeros_like(theta)\n",
        "  return tf.maximum(zero,\n",
        "                    -4.0 * tf.sin(theta - np.pi) * tf.cos(phi - 2.5) - 3.0)\n",
        "\n",
        "\n",
        "light_image_width = 30\n",
        "light_image_height = 30\n",
        "type = np.float64\n",
        "\n",
        "# Builds the pixels grid and compute corresponding spherical coordinates.\n",
        "spherical_coordinates = image_to_spherical_coordinates(light_image_width,\n",
        "                                                       light_image_height)\n",
        "theta = spherical_coordinates[:, :, 1]\n",
        "phi = spherical_coordinates[:, :, 2]\n",
        "\n",
        "# Samples the light function.\n",
        "sampled_light_function = light_function(theta, phi)\n",
        "ones_normal = tf.ones_like(theta)\n",
        "spherical_coordinates_3d = tf.stack((ones_normal, theta, phi), axis=-1)\n",
        "samples_direction_to_light = tf_math.spherical_to_cartesian_coordinates(\n",
        "    spherical_coordinates_3d)\n",
        "\n",
        "# Samples the SH.\n",
        "l, m = spherical_harmonics.generate_l_m_permutations(max_band)\n",
        "l = tf.convert_to_tensor(l)\n",
        "m = tf.convert_to_tensor(m)\n",
        "l_broadcasted = tf.broadcast_to(l, [light_image_width, light_image_height] +\n",
        "                                l.shape.as_list())\n",
        "m_broadcasted = tf.broadcast_to(m, [light_image_width, light_image_height] +\n",
        "                                l.shape.as_list())\n",
        "theta = tf.expand_dims(theta, axis=-1)\n",
        "theta_broadcasted = tf.broadcast_to(\n",
        "    theta, [light_image_width, light_image_height, 1])\n",
        "phi = tf.expand_dims(phi, axis=-1)\n",
        "phi_broadcasted = tf.broadcast_to(phi, [light_image_width, light_image_height, 1])\n",
        "sh_coefficients = spherical_harmonics.evaluate_spherical_harmonics(\n",
        "    l_broadcasted, m_broadcasted, theta_broadcasted, phi_broadcasted)\n",
        "sampled_light_function_broadcasted = tf.expand_dims(\n",
        "    sampled_light_function, axis=-1)\n",
        "sampled_light_function_broadcasted = tf.broadcast_to(\n",
        "    sampled_light_function_broadcasted,\n",
        "    [light_image_width, light_image_height] + l.shape.as_list())\n",
        "\n",
        "# Integrates the light function times SH over the sphere.\n",
        "projection = sh_coefficients * sampled_light_function_broadcasted * 4.0 * math.pi / (\n",
        "    light_image_width * light_image_height)\n",
        "light_coeffs = tf.reduce_sum(projection, (0, 1))\n",
        "\n",
        "# Reconstructs the image.\n",
        "reconstructed_light_function = tf.squeeze(\n",
        "    vector.dot(sh_coefficients, light_coeffs))\n",
        "\n",
        "print(\n",
        "    \"average l2 reconstruction error \",\n",
        "    np.linalg.norm(sampled_light_function - reconstructed_light_function) /\n",
        "    (light_image_width * light_image_height))\n",
        "\n",
        "vmin = np.minimum(\n",
        "    np.amin(np.minimum(sampled_light_function, reconstructed_light_function)),\n",
        "    0.0)\n",
        "vmax = np.maximum(\n",
        "    np.amax(np.maximum(sampled_light_function, reconstructed_light_function)),\n",
        "    1.0)\n",
        "# Plots results.\n",
        "plt.figure(figsize=(10, 10))\n",
        "ax = plt.subplot(\"131\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Original lighting function\")\n",
        "_ = ax.imshow(sampled_light_function, vmin=vmin, vmax=vmax)\n",
        "ax = plt.subplot(\"132\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Spherical Harmonics approximation\")\n",
        "_ = ax.imshow(reconstructed_light_function, vmin=vmin, vmax=vmax)\n",
        "ax = plt.subplot(\"133\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Difference\")\n",
        "_ = ax.imshow(\n",
        "    np.abs(reconstructed_light_function - sampled_light_function),\n",
        "    vmin=vmin,\n",
        "    vmax=vmax)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FReS5cj8UpPm"
      },
      "source": [
        "## Approximates the Lambertian BRDF with Zonal Harmonics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "n1yaUBTrac6C"
      },
      "outputs": [],
      "source": [
        "#################################################################\n",
        "# This cell creates an SH that approximates the Lambertian BRDF #\n",
        "#################################################################\n",
        "\n",
        "# The image dimensions control how many uniform samples we draw from the BRDF.\n",
        "brdf_image_width = 30\n",
        "brdf_image_height = 30\n",
        "type = np.float64\n",
        "\n",
        "# Builds the pixels grid and compute corresponding spherical coordinates.\n",
        "spherical_coordinates = image_to_spherical_coordinates(brdf_image_width,\n",
        "                                                       brdf_image_height)\n",
        "\n",
        "# Samples the BRDF function.\n",
        "cos_theta = tf.cos(spherical_coordinates[:, :, 1])\n",
        "sampled_brdf = tf.maximum(tf.zeros_like(cos_theta), cos_theta / np.pi)\n",
        "\n",
        "# Samples the zonal SH.\n",
        "l, m = spherical_harmonics.generate_l_m_zonal(max_band)\n",
        "l_broadcasted = tf.broadcast_to(l, [brdf_image_width, brdf_image_height] +\n",
        "                                l.shape.as_list())\n",
        "m_broadcasted = tf.broadcast_to(m, [brdf_image_width, brdf_image_height] +\n",
        "                                l.shape.as_list())\n",
        "theta = tf.expand_dims(spherical_coordinates[:, :, 1], axis=-1)\n",
        "theta_broadcasted = tf.broadcast_to(\n",
        "    theta, [brdf_image_width, brdf_image_height, 1])\n",
        "phi = tf.expand_dims(spherical_coordinates[:, :, 2], axis=-1)\n",
        "phi_broadcasted = tf.broadcast_to(phi, [brdf_image_width, brdf_image_height, 1])\n",
        "sh_coefficients = spherical_harmonics.evaluate_spherical_harmonics(\n",
        "    l_broadcasted, m_broadcasted, theta_broadcasted, phi_broadcasted)\n",
        "sampled_brdf_broadcasted = tf.expand_dims(sampled_brdf, axis=-1)\n",
        "sampled_brdf_broadcasted = tf.broadcast_to(\n",
        "    sampled_brdf_broadcasted,\n",
        "    [brdf_image_width, brdf_image_height] + l.shape.as_list())\n",
        "\n",
        "# Integrates the BRDF function times SH over the sphere.\n",
        "projection = sh_coefficients * sampled_brdf_broadcasted * 4.0 * math.pi / (\n",
        "    brdf_image_width * brdf_image_height)\n",
        "brdf_coeffs = tf.reduce_sum(projection, (0, 1))\n",
        "\n",
        "# Reconstructs the image.\n",
        "reconstructed_brdf = tf.squeeze(vector.dot(sh_coefficients, brdf_coeffs))\n",
        "\n",
        "print(\n",
        "    \"average l2 reconstruction error \",\n",
        "    np.linalg.norm(sampled_brdf - reconstructed_brdf) /\n",
        "    (brdf_image_width * brdf_image_height))\n",
        "\n",
        "vmin = np.minimum(np.amin(np.minimum(sampled_brdf, reconstructed_brdf)), 0.0)\n",
        "vmax = np.maximum(\n",
        "    np.amax(np.maximum(sampled_brdf, reconstructed_brdf)), 1.0 / np.pi)\n",
        "# Plots results.\n",
        "plt.figure(figsize=(10, 10))\n",
        "ax = plt.subplot(\"131\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Original reflectance function\")\n",
        "_ = ax.imshow(sampled_brdf, vmin=vmin, vmax=vmax)\n",
        "ax = plt.subplot(\"132\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Zonal Harmonics approximation\")\n",
        "_ = ax.imshow(reconstructed_brdf, vmin=vmin, vmax=vmax)\n",
        "ax = plt.subplot(\"133\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"Difference\")\n",
        "_ = ax.imshow(np.abs(sampled_brdf - reconstructed_brdf), vmin=vmin, vmax=vmax)\n",
        "\n",
        "plt.figure(figsize=(10, 5))\n",
        "plt.plot(\n",
        "    spherical_coordinates[:, 0, 1],\n",
        "    sampled_brdf[:, 0],\n",
        "    label=\"max(0,cos(x) / pi)\")\n",
        "plt.plot(\n",
        "    spherical_coordinates[:, 0, 1],\n",
        "    reconstructed_brdf[:, 0],\n",
        "    label=\"SH approximation\")\n",
        "plt.title(\"Approximation quality\")\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ra-Z2F6eUs3I"
      },
      "source": [
        "## Rotation of Zonal Harmonics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "s1kiXl8v0cin"
      },
      "outputs": [],
      "source": [
        "###############################\n",
        "# Rotation of zonal harmonics #\n",
        "###############################\n",
        "\n",
        "r_theta = tf.constant(np.pi / 2, shape=(1,), dtype=brdf_coeffs.dtype)\n",
        "r_phi = tf.constant(0.0, shape=(1,), dtype=brdf_coeffs.dtype)\n",
        "rotated_zonal_coefficients = spherical_harmonics.rotate_zonal_harmonics(\n",
        "    brdf_coeffs, r_theta, r_phi)\n",
        "\n",
        "# Builds the pixels grid and compute corresponding spherical coordinates.\n",
        "pixel_grid_start = np.array((0, 0), dtype=type)\n",
        "pixel_grid_end = np.array((brdf_image_width - 1, brdf_image_height - 1),\n",
        "                          dtype=type)\n",
        "pixel_nb = np.array((brdf_image_width, brdf_image_height))\n",
        "pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "normalized_pixels = pixels / (brdf_image_width - 1, brdf_image_height - 1)\n",
        "spherical_coordinates = tf_math.square_to_spherical_coordinates(\n",
        "    normalized_pixels)\n",
        "\n",
        "# reconstruction.\n",
        "l, m = spherical_harmonics.generate_l_m_permutations(max_band)\n",
        "l_broadcasted = tf.broadcast_to(\n",
        "    l, [light_image_width, light_image_height] + l.shape.as_list())\n",
        "m_broadcasted = tf.broadcast_to(\n",
        "    m, [light_image_width, light_image_height] + l.shape.as_list())\n",
        "theta = tf.expand_dims(spherical_coordinates[:, :, 1], axis=-1)\n",
        "theta_broadcasted = tf.broadcast_to(\n",
        "    theta, [light_image_width, light_image_height, 1])\n",
        "phi = tf.expand_dims(spherical_coordinates[:, :, 2], axis=-1)\n",
        "phi_broadcasted = tf.broadcast_to(\n",
        "    phi, [light_image_width, light_image_height, 1])\n",
        "sh_coefficients = spherical_harmonics.evaluate_spherical_harmonics(\n",
        "    l_broadcasted, m_broadcasted, theta_broadcasted, phi_broadcasted)\n",
        "\n",
        "reconstructed_rotated_brdf_function = tf.squeeze(\n",
        "    vector.dot(sh_coefficients, rotated_zonal_coefficients))\n",
        "\n",
        "plt.figure(figsize=(10, 10))\n",
        "ax = plt.subplot(\"121\")\n",
        "ax.set_title(\"Zonal SH\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "_ = ax.imshow(reconstructed_brdf)\n",
        "ax = plt.subplot(\"122\")\n",
        "ax.set_title(\"Rotated version\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "_ = ax.imshow(reconstructed_rotated_brdf_function)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ha31toe2UycZ"
      },
      "source": [
        "## Reconstruction via Spherical Harmonics convolution of the SH lighting and ZH BRDF"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "I31bzJ58YjJa"
      },
      "outputs": [],
      "source": [
        "############################################################################################\n",
        "# Helper function allowing to estimate sphere normal and depth for each pixel in the image #\n",
        "############################################################################################\n",
        "def compute_intersection_normal_sphere(image_width, image_height, sphere_radius,\n",
        "                                       sphere_center, type):\n",
        "  pixel_grid_start = np.array((0.5, 0.5), dtype=type)\n",
        "  pixel_grid_end = np.array((image_width - 0.5, image_height - 0.5), dtype=type)\n",
        "  pixel_nb = np.array((image_width, image_height))\n",
        "  pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "\n",
        "  pixel_ray = tf.math.l2_normalize(orthographic.ray(pixels), axis=-1)\n",
        "  zero_depth = np.zeros([image_width, image_height, 1])\n",
        "  pixels_3d = orthographic.unproject(pixels, zero_depth)\n",
        "\n",
        "  intersections_points, normals = ray.intersection_ray_sphere(\n",
        "      sphere_center, sphere_radius, pixel_ray, pixels_3d)\n",
        "  return intersections_points[0, :, :, :], normals[0, :, :, :]\n",
        "\n",
        "\n",
        "###############################\n",
        "# Setup the image, and sphere #\n",
        "###############################\n",
        "# Image dimensions\n",
        "image_width = 100\n",
        "image_height = 80\n",
        "\n",
        "# Sphere center and radius\n",
        "sphere_radius = np.array((30.0,), dtype=type)\n",
        "sphere_center = np.array((image_width / 2.0, image_height / 2.0, 100.0),\n",
        "                         dtype=type)\n",
        "\n",
        "# Builds the pixels grid and compute corresponding spherical coordinates.\n",
        "pixel_grid_start = np.array((0, 0), dtype=type)\n",
        "pixel_grid_end = np.array((image_width - 1, image_height - 1), dtype=type)\n",
        "pixel_nb = np.array((image_width, image_height))\n",
        "pixels = grid.generate(pixel_grid_start, pixel_grid_end, pixel_nb)\n",
        "normalized_pixels = pixels / (image_width - 1, image_height - 1)\n",
        "spherical_coordinates = tf_math.square_to_spherical_coordinates(\n",
        "    normalized_pixels)\n",
        "\n",
        "################################################################################################\n",
        "# For each pixel in the image, estimate the corresponding surface point and associated normal. #\n",
        "################################################################################################\n",
        "intersection_3d, surface_normal = compute_intersection_normal_sphere(\n",
        "    image_width, image_height, sphere_radius, sphere_center, type)\n",
        "\n",
        "surface_normals_spherical_coordinates = tf_math.cartesian_to_spherical_coordinates(\n",
        "    surface_normal)\n",
        "\n",
        "# SH\n",
        "l, m = spherical_harmonics.generate_l_m_permutations(\n",
        "    max_band)  # recomputed =\u003e optimize\n",
        "l = tf.convert_to_tensor(l)\n",
        "m = tf.convert_to_tensor(m)\n",
        "l_broadcasted = tf.broadcast_to(l,\n",
        "                                [image_width, image_height] + l.shape.as_list())\n",
        "m_broadcasted = tf.broadcast_to(m,\n",
        "                                [image_width, image_height] + l.shape.as_list())\n",
        "\n",
        "#################################################\n",
        "# Estimates result using SH convolution - cheap #\n",
        "#################################################\n",
        "\n",
        "sh_integration = spherical_harmonics.integration_product(\n",
        "    light_coeffs,\n",
        "    spherical_harmonics.rotate_zonal_harmonics(\n",
        "        brdf_coeffs,\n",
        "        tf.expand_dims(surface_normals_spherical_coordinates[:, :, 1], axis=-1),\n",
        "        tf.expand_dims(surface_normals_spherical_coordinates[:, :, 2],\n",
        "                       axis=-1)),\n",
        "    keepdims=False)\n",
        "\n",
        "# Sets pixels not belonging to the sphere to 0.\n",
        "sh_integration = tf.where(\n",
        "    tf.greater(intersection_3d[:, :, 2], 0.0), sh_integration,\n",
        "    tf.zeros_like(sh_integration))\n",
        "# Sets pixels with negative light to 0.\n",
        "sh_integration = tf.where(\n",
        "    tf.greater(sh_integration, 0.0), sh_integration,\n",
        "    tf.zeros_like(sh_integration))\n",
        "\n",
        "###########################################\n",
        "# 'Brute force' solution - very expensive #\n",
        "###########################################\n",
        "\n",
        "factor = 4.0 * np.pi / (light_image_width * light_image_height)\n",
        "gt = tf.einsum(\n",
        "    \"hwn,uvn-\u003ehwuv\", surface_normal,\n",
        "    samples_direction_to_light *\n",
        "    tf.expand_dims(sampled_light_function, axis=-1))\n",
        "gt = tf.maximum(gt, 0.0)  # removes negative dot products\n",
        "gt = tf.reduce_sum(gt, axis=(2, 3))\n",
        "# Sets pixels not belonging to the sphere to 0.\n",
        "gt = tf.where(tf.greater(intersection_3d[:, :, 2], 0.0), gt, tf.zeros_like(gt))\n",
        "gt *= factor\n",
        "\n",
        "# TODO(b/124463095): gt and sh_integration differ by a factor of pi.\n",
        "sh_integration = np.transpose(sh_integration, (1, 0))\n",
        "gt = np.transpose(gt, (1, 0))\n",
        "\n",
        "plt.figure(figsize=(10, 20))\n",
        "ax = plt.subplot(\"121\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"SH light and SH BRDF\")\n",
        "_ = ax.imshow(sh_integration, vmin=0.0)\n",
        "ax = plt.subplot(\"122\")\n",
        "ax.axes.get_xaxis().set_visible(False)\n",
        "ax.axes.get_yaxis().set_visible(False)\n",
        "ax.grid(False)\n",
        "ax.set_title(\"GT light and GT BRDF\")\n",
        "_ = ax.imshow(gt, vmin=0.0)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/brain/python/client:colab_notebook",
        "kind": "private"
      },
      "name": "spherical_harmonics_demo",
      "provenance": [],
      "toc_visible": true,
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
